import whisper
import opencc
import sys
import numpy as np

def transcribe(audio_file):
    # 加载 Whisper 模型
    model = whisper.load_model("base")
    
    # 转录音频
    result = model.transcribe(audio_file, word_timestamps=True)

    print("result===", result)
    
    # 获取转录文本和时间戳信息
    transcribed_text = result['text']

    # 初始化 opencc 转换器，转换繁体到简体
    cc = opencc.OpenCC('t2s.json')  # 繁体转简体，t2s.json 是转换配置文件
    transcribed_text_simplified = cc.convert(transcribed_text)  # 转换文本
    
    # 获取时间信息
    segments = result['segments']  # 获取每个段落的时间戳信息
    
    time_info = []
    for segment in segments:
        start_time = segment['start']  # 开始时间
        end_time = segment['end']  # 结束时间
        text = segment['text']  # 转录的文本
        text_simplified = cc.convert(text)  # 转换每个段落的文本
        time_info.append({
            # 'start': {start_time},
            # 'end': np.round(end_time, 2),
            'start': f"{np.round(start_time, 2)}",
            'end': f"{np.round(end_time, 2)}",
            'text': text_simplified
        })
    
    return transcribed_text_simplified, time_info

if __name__ == "__main__":
    audio_file = sys.argv[1]  # 获取命令行传入的音频文件路径
    transcribed_text, time_info = transcribe(audio_file)
    
    # 打印转录文本（简体）
    print("Transcribed Text (Simplified):")
    print(transcribed_text)
    print(time_info) 
    # 打印时间信息
    print("\nTime Information:")
    for item in time_info:
	   # print(item)
        print(f"Start: {item['start']}s, End: {item['end']}s, Text: {item['text']}")

